{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/tensorflow-without-a-phd/blob/master/tensorflow-mnist-tutorial/keras_04_mnist_convolutional.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Lvo0t7XVIkWZ"
   },
   "source": [
    "### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cCpkS9C_H7Tl"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "EPOCHS = 10\n",
    "\n",
    "training_images_file   = 'gs://mnist-public/train-images-idx3-ubyte'\n",
    "training_labels_file   = 'gs://mnist-public/train-labels-idx1-ubyte'\n",
    "validation_images_file = 'gs://mnist-public/t10k-images-idx3-ubyte'\n",
    "validation_labels_file = 'gs://mnist-public/t10k-labels-idx1-ubyte'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qpiJj8ym0v0-"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "AoilhmYe1b5t",
    "outputId": "167b0893-ab0b-4179-b28b-6a5edc480956"
   },
   "outputs": [],
   "source": [
    "import os, re, math, json, shutil, pprint\n",
    "import PIL.Image, PIL.ImageFont, PIL.ImageDraw\n",
    "import IPython.display as display\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from matplotlib import pyplot as plt\n",
    "print(\"Tensorflow version \" + tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "qhdz68Xm3Z4Z"
   },
   "outputs": [],
   "source": [
    "#@title visualization utilities [RUN ME]\n",
    "\"\"\"\n",
    "This cell contains helper functions used for visualization\n",
    "and downloads only. You can skip reading it. There is very\n",
    "little useful Keras/Tensorflow code here.\n",
    "\"\"\"\n",
    "\n",
    "# Matplotlib config\n",
    "plt.ioff()\n",
    "plt.rc('image', cmap='gray_r')\n",
    "plt.rc('grid', linewidth=1)\n",
    "plt.rc('xtick', top=False, bottom=False, labelsize='large')\n",
    "plt.rc('ytick', left=False, right=False, labelsize='large')\n",
    "plt.rc('axes', facecolor='F8F8F8', titlesize=\"large\", edgecolor='white')\n",
    "plt.rc('text', color='a8151a')\n",
    "plt.rc('figure', facecolor='F0F0F0', figsize=(16,9))\n",
    "# Matplotlib fonts\n",
    "MATPLOTLIB_FONT_DIR = os.path.join(os.path.dirname(plt.__file__), \"mpl-data/fonts/ttf\")\n",
    "\n",
    "# pull a batch from the datasets. This code is not very nice, it gets much better in eager mode (TODO)\n",
    "def dataset_to_numpy_util(training_dataset, validation_dataset, N):\n",
    "  \n",
    "  # get one batch from each: 10000 validation digits, N training digits\n",
    "  batch_train_ds = training_dataset.unbatch().batch(N)\n",
    "  \n",
    "  # eager execution: loop through datasets normally\n",
    "  if tf.executing_eagerly():\n",
    "    for validation_digits, validation_labels in validation_dataset:\n",
    "      validation_digits = validation_digits.numpy()\n",
    "      validation_labels = validation_labels.numpy()\n",
    "      break\n",
    "    for training_digits, training_labels in batch_train_ds:\n",
    "      training_digits = training_digits.numpy()\n",
    "      training_labels = training_labels.numpy()\n",
    "      break\n",
    "    \n",
    "  else:\n",
    "    v_images, v_labels = validation_dataset.make_one_shot_iterator().get_next()\n",
    "    t_images, t_labels = batch_train_ds.make_one_shot_iterator().get_next()\n",
    "    # Run once, get one batch. Session.run returns numpy results\n",
    "    with tf.Session() as ses:\n",
    "      (validation_digits, validation_labels,\n",
    "       training_digits, training_labels) = ses.run([v_images, v_labels, t_images, t_labels])\n",
    "  \n",
    "  # these were one-hot encoded in the dataset\n",
    "  validation_labels = np.argmax(validation_labels, axis=1)\n",
    "  training_labels = np.argmax(training_labels, axis=1)\n",
    "  \n",
    "  return (training_digits, training_labels,\n",
    "          validation_digits, validation_labels)\n",
    "\n",
    "# create digits from local fonts for testing\n",
    "def create_digits_from_local_fonts(n):\n",
    "  font_labels = []\n",
    "  img = PIL.Image.new('LA', (28*n, 28), color = (0,255)) # format 'LA': black in channel 0, alpha in channel 1\n",
    "  font1 = PIL.ImageFont.truetype(os.path.join(MATPLOTLIB_FONT_DIR, 'DejaVuSansMono-Oblique.ttf'), 25)\n",
    "  font2 = PIL.ImageFont.truetype(os.path.join(MATPLOTLIB_FONT_DIR, 'STIXGeneral.ttf'), 25)\n",
    "  d = PIL.ImageDraw.Draw(img)\n",
    "  for i in range(n):\n",
    "    font_labels.append(i%10)\n",
    "    d.text((7+i*28,0 if i<10 else -4), str(i%10), fill=(255,255), font=font1 if i<10 else font2)\n",
    "  font_digits = np.array(img.getdata(), np.float32)[:,0] / 255.0 # black in channel 0, alpha in channel 1 (discarded)\n",
    "  font_digits = np.reshape(np.stack(np.split(np.reshape(font_digits, [28, 28*n]), n, axis=1), axis=0), [n, 28*28])\n",
    "  return font_digits, font_labels\n",
    "\n",
    "# utility to display a row of digits with their predictions\n",
    "def display_digits(digits, predictions, labels, title, n):\n",
    "  fig = plt.figure(figsize=(13,3))\n",
    "  digits = np.reshape(digits, [n, 28, 28])\n",
    "  digits = np.swapaxes(digits, 0, 1)\n",
    "  digits = np.reshape(digits, [28, 28*n])\n",
    "  plt.yticks([])\n",
    "  plt.xticks([28*x+14 for x in range(n)], predictions)\n",
    "  plt.grid(b=None)\n",
    "  for i,t in enumerate(plt.gca().xaxis.get_ticklabels()):\n",
    "    if predictions[i] != labels[i]: t.set_color('red') # bad predictions in red\n",
    "  plt.imshow(digits)\n",
    "  plt.grid(None)\n",
    "  plt.title(title)\n",
    "  display.display(fig)\n",
    "  \n",
    "# utility to display multiple rows of digits, sorted by unrecognized/recognized status\n",
    "def display_top_unrecognized(digits, predictions, labels, n, lines):\n",
    "  idx = np.argsort(predictions==labels) # sort order: unrecognized first\n",
    "  for i in range(lines):\n",
    "    display_digits(digits[idx][i*n:(i+1)*n], predictions[idx][i*n:(i+1)*n], labels[idx][i*n:(i+1)*n],\n",
    "                   \"{} sample validation digits out of {} with bad predictions in red and sorted first\".format(n*lines, len(digits)) if i==0 else \"\", n)\n",
    "\n",
    "def plot_learning_rate(lr_func, epochs):\n",
    "  xx = np.arange(epochs+1, dtype=np.float)\n",
    "  y = [lr_decay(x) for x in xx]\n",
    "  fig, ax = plt.subplots(figsize=(9, 6))\n",
    "  ax.set_xlabel('epochs')\n",
    "  ax.set_title('Learning rate\\ndecays from {:0.3g} to {:0.3g}'.format(y[0], y[-2]))\n",
    "  ax.minorticks_on()\n",
    "  ax.grid(True, which='major', axis='both', linestyle='-', linewidth=1)\n",
    "  ax.grid(True, which='minor', axis='both', linestyle=':', linewidth=0.5)\n",
    "  ax.step(xx,y, linewidth=3, where='post')\n",
    "  display.display(fig)\n",
    "\n",
    "class PlotTraining(tf.keras.callbacks.Callback):\n",
    "  def __init__(self, sample_rate=1, zoom=1):\n",
    "    self.sample_rate = sample_rate\n",
    "    self.step = 0\n",
    "    self.zoom = zoom\n",
    "    self.steps_per_epoch = 60000//BATCH_SIZE\n",
    "\n",
    "  def on_train_begin(self, logs={}):\n",
    "    self.batch_history = {}\n",
    "    self.batch_step = []\n",
    "    self.epoch_history = {}\n",
    "    self.epoch_step = []\n",
    "    self.fig, self.axes = plt.subplots(1, 2, figsize=(16, 7))\n",
    "    plt.ioff()\n",
    "\n",
    "  def on_batch_end(self, batch, logs={}):\n",
    "    if (batch % self.sample_rate) == 0:\n",
    "      self.batch_step.append(self.step)\n",
    "      for k,v in logs.items():\n",
    "        # do not log \"batch\" and \"size\" metrics that do not change\n",
    "        # do not log training accuracy \"acc\"\n",
    "        if k=='batch' or k=='size':# or k=='acc':\n",
    "          continue\n",
    "        self.batch_history.setdefault(k, []).append(v)\n",
    "    self.step += 1\n",
    "\n",
    "  def on_epoch_end(self, epoch, logs={}):\n",
    "    plt.close(self.fig)\n",
    "    self.axes[0].cla()\n",
    "    self.axes[1].cla()\n",
    "      \n",
    "    self.axes[0].set_ylim(0, 1.2/self.zoom)\n",
    "    self.axes[1].set_ylim(1-1/self.zoom/2, 1+0.1/self.zoom/2)\n",
    "    \n",
    "    self.epoch_step.append(self.step)\n",
    "    for k,v in logs.items():\n",
    "      # only log validation metrics\n",
    "      if not k.startswith('val_'):\n",
    "        continue\n",
    "      self.epoch_history.setdefault(k, []).append(v)\n",
    "\n",
    "    display.clear_output(wait=True)\n",
    "    \n",
    "    for k,v in self.batch_history.items():\n",
    "      self.axes[0 if k.endswith('loss') else 1].plot(np.array(self.batch_step) / self.steps_per_epoch, v, label=k)\n",
    "      \n",
    "    for k,v in self.epoch_history.items():\n",
    "      self.axes[0 if k.endswith('loss') else 1].plot(np.array(self.epoch_step) / self.steps_per_epoch, v, label=k, linewidth=3)\n",
    "      \n",
    "    self.axes[0].legend()\n",
    "    self.axes[1].legend()\n",
    "    self.axes[0].set_xlabel('epochs')\n",
    "    self.axes[1].set_xlabel('epochs')\n",
    "    self.axes[0].minorticks_on()\n",
    "    self.axes[0].grid(True, which='major', axis='both', linestyle='-', linewidth=1)\n",
    "    self.axes[0].grid(True, which='minor', axis='both', linestyle=':', linewidth=0.5)\n",
    "    self.axes[1].minorticks_on()\n",
    "    self.axes[1].grid(True, which='major', axis='both', linestyle='-', linewidth=1)\n",
    "    self.axes[1].grid(True, which='minor', axis='both', linestyle=':', linewidth=0.5)\n",
    "    display.display(self.fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Lz1Zknfk4qCx"
   },
   "source": [
    "### tf.data.Dataset: parse files and prepare training and validation datasets\n",
    "Please read the [best practices for building](https://www.tensorflow.org/guide/performance/datasets) input pipelines with tf.data.Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZE8dgyPC1_6m"
   },
   "outputs": [],
   "source": [
    "AUTO = tf.data.experimental.AUTOTUNE\n",
    "\n",
    "def read_label(tf_bytestring):\n",
    "    label = tf.io.decode_raw(tf_bytestring, tf.uint8)\n",
    "    label = tf.reshape(label, [])\n",
    "    label = tf.one_hot(label, 10)\n",
    "    return label\n",
    "  \n",
    "def read_image(tf_bytestring):\n",
    "    image = tf.io.decode_raw(tf_bytestring, tf.uint8)\n",
    "    image = tf.cast(image, tf.float32)/256.0\n",
    "    image = tf.reshape(image, [28*28])\n",
    "    return image\n",
    "  \n",
    "def load_dataset(image_file, label_file):\n",
    "    imagedataset = tf.data.FixedLengthRecordDataset(image_file, 28*28, header_bytes=16)\n",
    "    imagedataset = imagedataset.map(read_image, num_parallel_calls=16)\n",
    "    labelsdataset = tf.data.FixedLengthRecordDataset(label_file, 1, header_bytes=8)\n",
    "    labelsdataset = labelsdataset.map(read_label, num_parallel_calls=16)\n",
    "    dataset = tf.data.Dataset.zip((imagedataset, labelsdataset))\n",
    "    return dataset \n",
    "  \n",
    "def get_training_dataset(image_file, label_file, batch_size):\n",
    "    dataset = load_dataset(image_file, label_file)\n",
    "    dataset = dataset.cache()  # this small dataset can be entirely cached in RAM, for TPU this is important to get good performance from such a small dataset\n",
    "    dataset = dataset.shuffle(5000, reshuffle_each_iteration=True)\n",
    "    dataset = dataset.repeat() # Mandatory for Keras for now\n",
    "    dataset = dataset.batch(batch_size, drop_remainder=True) # drop_remainder is important on TPU, batch size must be fixed\n",
    "    dataset = dataset.prefetch(AUTO)  # fetch next batches while training on the current one (-1: autotune prefetch buffer size)\n",
    "    return dataset\n",
    "  \n",
    "def get_validation_dataset(image_file, label_file):\n",
    "    dataset = load_dataset(image_file, label_file)\n",
    "    dataset = dataset.cache() # this small dataset can be entirely cached in RAM, for TPU this is important to get good performance from such a small dataset\n",
    "    dataset = dataset.batch(10000, drop_remainder=True) # 10000 items in eval dataset, all in one batch\n",
    "    dataset = dataset.repeat() # Mandatory for Keras for now\n",
    "    return dataset\n",
    "\n",
    "# instantiate the datasets\n",
    "training_dataset = get_training_dataset(training_images_file, training_labels_file, BATCH_SIZE)\n",
    "validation_dataset = get_validation_dataset(validation_images_file, validation_labels_file)\n",
    "\n",
    "# For TPU, we will need a function that returns the dataset\n",
    "training_input_fn = lambda: get_training_dataset(training_images_file, training_labels_file, BATCH_SIZE)\n",
    "validation_input_fn = lambda: get_validation_dataset(validation_images_file, validation_labels_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_fXo6GuvL3EB"
   },
   "source": [
    "### Let's have a look at the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 177
    },
    "colab_type": "code",
    "id": "yZ4tjPKvL2eh",
    "outputId": "f69b308a-3ceb-445d-aeee-1fcef36b3bb5"
   },
   "outputs": [],
   "source": [
    "N = 24\n",
    "(training_digits, training_labels,\n",
    " validation_digits, validation_labels) = dataset_to_numpy_util(training_dataset, validation_dataset, N)\n",
    "display_digits(training_digits, training_labels, training_labels, \"training digits and their labels\", N)\n",
    "display_digits(validation_digits[:N], validation_labels[:N], validation_labels[:N], \"validation digits and their labels\", N)\n",
    "font_digits, font_labels = create_digits_from_local_fonts(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KIc0oqiD40HC"
   },
   "source": [
    "### Keras model\n",
    "If you are not sure what cross-entropy, dropout, softmax or batch-normalization mean, head here for a crash-course: [Tensorflow and deep learning without a PhD](https://github.com/GoogleCloudPlatform/tensorflow-without-a-phd/#featured-code-sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 391
    },
    "colab_type": "code",
    "id": "56y8UNFQIVwj",
    "outputId": "e2fa4d89-ec5b-4508-950e-31ad0b87f09e"
   },
   "outputs": [],
   "source": [
    "model = tf.keras.Sequential(\n",
    "  [\n",
    "      tf.keras.layers.Reshape(input_shape=(28*28,), target_shape=(28, 28, 1)),\n",
    "      tf.keras.layers.Conv2D(kernel_size=3, filters=12, activation='relu', padding='same'),\n",
    "      tf.keras.layers.Conv2D(kernel_size=6, filters=24, activation='relu', padding='same', strides=2),\n",
    "      tf.keras.layers.Conv2D(kernel_size=6, filters=32, activation='relu', padding='same', strides=2),\n",
    "      tf.keras.layers.Flatten(),\n",
    "      tf.keras.layers.Dense(200, activation='relu'),\n",
    "      tf.keras.layers.Dense(10, activation='softmax')\n",
    "  ])\n",
    "\n",
    "model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.01),\n",
    "              loss='categorical_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "# print model layers\n",
    "model.summary()\n",
    "\n",
    "# utility callback that displays training curves\n",
    "plot_training = PlotTraining(sample_rate=10, zoom=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "E0A7CSbW67EY"
   },
   "source": [
    "### Learning Rate schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 422
    },
    "colab_type": "code",
    "id": "2qz_k44U6_-f",
    "outputId": "e996d4ec-36da-4f18-cb20-27d2ed451668"
   },
   "outputs": [],
   "source": [
    "# lr decay function\n",
    "def lr_decay(epoch):\n",
    "  return 0.01 * math.pow(0.6, epoch)\n",
    "\n",
    "# lr schedule callback\n",
    "lr_decay_callback = tf.keras.callbacks.LearningRateScheduler(lr_decay, verbose=True)\n",
    "\n",
    "# important to see what you are doing\n",
    "plot_learning_rate(lr_decay, EPOCHS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CuhDh8ao8VyB"
   },
   "source": [
    "### Train and validate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 466
    },
    "colab_type": "code",
    "id": "TTwH_P-ZJ_xx",
    "outputId": "e99eae66-946b-433f-bef1-886617b573d5"
   },
   "outputs": [],
   "source": [
    "steps_per_epoch = 60000//BATCH_SIZE  # 60,000 items in this dataset\n",
    "print(\"Steps per epoch: \", steps_per_epoch)\n",
    "\n",
    "history = model.fit(training_dataset, steps_per_epoch=steps_per_epoch, epochs=EPOCHS,\n",
    "                    validation_data=validation_dataset, validation_steps=1, callbacks=[plot_training, lr_decay_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9jFVovcUUVs1"
   },
   "source": [
    "### Visualize predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 585
    },
    "colab_type": "code",
    "id": "w12OId8Mz7dF",
    "outputId": "7575d8a9-f707-4448-c161-617e93766c15"
   },
   "outputs": [],
   "source": [
    "# recognize digits from local fonts\n",
    "probabilities = model.predict(font_digits, steps=1)\n",
    "predicted_labels = np.argmax(probabilities, axis=1)\n",
    "display_digits(font_digits, predicted_labels, font_labels, \"predictions from local fonts (bad predictions in red)\", N)\n",
    "\n",
    "# recognize validation digits\n",
    "probabilities = model.predict(validation_digits, steps=1)\n",
    "predicted_labels = np.argmax(probabilities, axis=1)\n",
    "display_top_unrecognized(validation_digits, predicted_labels, validation_labels, N, 7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SVY1pBg5ydH-"
   },
   "source": [
    "## License"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hleIN5-pcr0N"
   },
   "source": [
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "author: Martin Gorner<br>\n",
    "twitter: @martin_gorner\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "Copyright 2019 Google LLC\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "This is not an official Google product but sample code provided for an educational purpose\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "keras_04_mnist_convolutional.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "environment": {
   "name": "tf22-gpu.2-2.m47",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf22-gpu.2-2:m47"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
